#include "LearningAndRecognition.h"
#include "FacialRecognition.h"
#include <stdio.h>
#include <string.h>

const char * TRAIN_FACES = "train_faces.txt";
const char * RECOGNIZE_FACES = "recognize_face.txt";

//// Global variables
int nTrainFaces = 0; // number of training images
int nEigens = 0; // number of eigenvalues
IplImage ** faceImgArr = 0; // array of face images
CvMat * personNumTruthMat = 0; //array of person numbers
IplImage * pAvgTrainImg = 0; // the average image
IplImage ** eigenVectArr = 0; // eigenvectors
CvMat * eigenValMat = 0; // eigenvalues
CvMat * projectedTrainFaceMat = 0; // projected training faces
double leastDistSq = DBL_MAX; //Least distance to image
const double RECOGNITION_THRESHOLD = 0.20;
int recognitionCounter = 0;

LearningAndRecognition::LearningAndRecognition(void)
{
}

void LearningAndRecognition::LearnAndRecognizeFace(char * argv[])
{
	std::string optionInput = "";
	std::string input = "";
	std::vector<std::string> vectorFilePaths;
	std::vector<std::string> vectorPersonNames;
	int personCounter = 0;

	std::cout << "\n1. Train my face.\n";
	getline(std::cin, optionInput);

	if(optionInput == "1") {
		std::cout << "\nWhat is your name?\n";
		std::getline(std::cin, input);
		std::wstring stemp = s2ws(input); //Temporary buffer is required
		LPCWSTR dir = stemp.c_str();
		CreateDirectory(dir, NULL);
		CreateTrainingSetAndStoreOnDisk(argv, input, vectorFilePaths, vectorPersonNames, personCounter);

		if(CapturingAndStoring::killThread) {
			std::cout << "\nThe square has to be green! Please exit and try again...\n";
			exit(1);			
			std::terminate();
		}

		CreateTrainingTextFileWithFilePaths(vectorFilePaths);
		learn(argv);
		std::cout << "\nYour face is now trained, " << input << "! Please any key to exit...\n";
		std::getline(std::cin, input);
	}
}

std::wstring LearningAndRecognition::s2ws(const std::string& s)
{
	int len;
	int slength = (int)s.length() + 1;
	len = MultiByteToWideChar(CP_ACP, 0, s.c_str(), slength, 0, 0); 
	wchar_t* buf = new wchar_t[len];
	MultiByteToWideChar(CP_ACP, 0, s.c_str(), slength, buf, len);
	std::wstring r(buf);
	delete[] buf;
	return r;
}

LearningAndRecognition::~LearningAndRecognition(void)
{
}

void LearningAndRecognition::CreateTrainingSetAndStoreOnDisk(char * argv[], std::string input, std::vector<std::string> &vectorFilePaths, std::vector<std::string> &vectorPersonNames, int &personCounter)
{
	CapturingAndStoring capturingAndStoring;
	std::string dynInput = "";
	const char * filePath = "";
	bool personNameExists = false;
	personCounter++;
	std::vector<std::string>::iterator iteratorVector;
	iteratorVector = std::find(vectorPersonNames.begin(), vectorPersonNames.end(), input);
	
	if(iteratorVector == vectorPersonNames.end()) {
		personNameExists = false;
		vectorPersonNames.push_back(input);		
	}
	else
		personNameExists = true;

	//Prepare face for capturingAndStoring and store on disk
	capturingAndStoring.PrepareFaceInImageAndStoreOnDisk(argv, FacialRecognition::inputImage, FacialRecognition::faceRect, input, personCounter, dynInput, filePath, personNameExists, vectorFilePaths);
}

void LearningAndRecognition::CreateRecognitionTextFileWithFilePaths(std::vector<std::string> &vectorFilePaths)
{
	std::ofstream file;
	file.open(RECOGNIZE_FACES);
	file << vectorFilePaths[0] + "\n";
	file.close();
	vectorFilePaths.clear();
}

void LearningAndRecognition::CreateTrainingTextFileWithFilePaths(std::vector<std::string> &vectorFilePaths)
{
	std::fstream file;
	std::string filePath = "";
	std::string tempFilePath = "";
	char temp = NULL;
	int num = 0;
	file.open(TRAIN_FACES);

	while(!file.eof()) {
		getline(file, tempFilePath);
		if(tempFilePath.empty())
			break;
		else
			filePath = tempFilePath;
	}

	if(!filePath.empty()) {
		file.close();
		file.open(TRAIN_FACES, std::fstream::in | std::fstream::out | std::fstream::app);

		std::string tempNum = "";
		for(int i = 0; i < filePath.length(); i++) {
			if(filePath[i] == ' ')
				break;

			tempNum += filePath[i];				
		}
		num = atoi(tempNum.c_str());
		tempNum = "";
		num++;
		std::stringstream out;
		out << num;
		std::string stringNum = out.str();
		std::string tempFilePath = vectorFilePaths[0];

		for(int i = 0; i < tempFilePath.length(); i++) {
			if(tempFilePath[0] == ' ')
				break;
			tempFilePath.erase(0, 1);
		}

		tempFilePath = stringNum + tempFilePath;
		vectorFilePaths[0] = tempFilePath;
		file << vectorFilePaths[0] + "\n";
		file.close();
	}
	else {
		file.close();
		std::ofstream file;
		file.open(TRAIN_FACES);
		file << vectorFilePaths[0] + "\n";
		file.close();
	}

	vectorFilePaths.clear();
}

void LearningAndRecognition::ReadTrainingTextFile(std::map<int, std::string> &mapPersonNumName)
{
	std::fstream file;
	std::string filePath = "";
	int personNum  = 0;
	std::string name = "";
	file.open(TRAIN_FACES);

	while(!file.eof()) {
		getline(file, filePath);
		if(filePath.empty() && mapPersonNumName.empty()) {
			std::cout << "\nTraining text file is empty! Exiting...\n";
			exit(1);			
			std::terminate();
		}
		else {
			if(filePath.empty())
				break;
			//Store the person number and name
			int i = 0;
			std::string tempNum = "";
			for(i; i < filePath.length(); i++) {
				if(filePath[i] == ' ') {
					i++;
					break;
				}
				tempNum += filePath[i];				
			}
			personNum = atoi(tempNum.c_str());
			tempNum = "";
			for(i; i < filePath.length(); i++) {
				if(filePath[i] == '\\')
					break;
				name += filePath[i];
			}
			mapPersonNumName[personNum] = name;
			name = "";
		}		
	}
	
	file.close();
}

bool LearningAndRecognition::fexists(const char * FILE)
{
	std::ifstream file(FILE);
	return file;
}

void LearningAndRecognition::learn(char * argv[])
{
	int i;
	// load training data
	nTrainFaces = loadFaceImgArray(TRAIN_FACES);

	// do PCA on the training faces
	doPCA();

	// project the training images onto the PCA subspace
	projectedTrainFaceMat = cvCreateMat(nTrainFaces, nEigens, CV_32FC1);
	for(i=0; i<nTrainFaces; i++)
	{
		cvEigenDecomposite(
			faceImgArr[i],
			nEigens,
			eigenVectArr,
			0, 0,
			pAvgTrainImg,
			projectedTrainFaceMat->data.fl + i*nEigens);
	}
	// store the capturingAndStoring data as an xml file
	storeTrainingData(argv);
}

int LearningAndRecognition::loadFaceImgArray(const char * filename)
{
	FILE * imgListFile = 0;
	char imgFilename[512];
	int iFace, nFaces=0;

	// open the input file
	imgListFile = fopen(filename, "r");

	// count the number of faces
	while( fgets(imgFilename, 512, imgListFile) ) ++nFaces;
	rewind(imgListFile);

	// allocate the face-image array and person number matrix
	faceImgArr = (IplImage **)cvAlloc( nFaces*sizeof(IplImage *) );
	personNumTruthMat = cvCreateMat( 1, nFaces, CV_32SC1 );

	// store the face images in an array
	for(iFace=0; iFace<nFaces; iFace++)
	{
		// read person number and name of image file
		fscanf(imgListFile,
			"%d %s", personNumTruthMat->data.i+iFace, imgFilename);

		// load the face image
		faceImgArr[iFace] = cvLoadImage(imgFilename, CV_LOAD_IMAGE_GRAYSCALE);
	}

	fclose(imgListFile);

	return nFaces;
}

void LearningAndRecognition::doPCA()
{
	int i;
	CvTermCriteria calcLimit;
	CvSize faceImgSize;

	// set the number of eigenvalues to use
	nEigens = nTrainFaces-1;

	// allocate the eigenvector images
	faceImgSize.width = faceImgArr[0]->width;
	faceImgSize.height = faceImgArr[0]->height;
	eigenVectArr = (IplImage**)cvAlloc(sizeof(IplImage*) * nEigens);
	for(i=0; i<nEigens; i++)
		eigenVectArr[i] = cvCreateImage(faceImgSize, IPL_DEPTH_32F, 1);

	// allocate the eigenvalue array
	eigenValMat = cvCreateMat( 1, nEigens, CV_32FC1 );

	// allocate the averaged image
	pAvgTrainImg = cvCreateImage(faceImgSize, IPL_DEPTH_32F, 1);

	// set the PCA termination criterion
	calcLimit = cvTermCriteria( CV_TERMCRIT_ITER, nEigens, 1);

	// compute average image, eigenvalues, and eigenvectors
	cvCalcEigenObjects(nTrainFaces,
		(void*)faceImgArr,
		(void*)eigenVectArr,
		CV_EIGOBJ_NO_CALLBACK,
		0,
		0,
		&calcLimit,
		pAvgTrainImg,
		eigenValMat->data.fl);
}

void LearningAndRecognition::storeTrainingData(char * argv[])
{
	CvFileStorage * fileStorage;
	int i;

	std::stringstream ss;
	std::string facesDataString;
	ss << argv[1];
	ss >> facesDataString;
	facesDataString += "\\facesdata.xml";
	const char * FACES_DATA = facesDataString.c_str();

	// create a file-storage interface
	fileStorage = cvOpenFileStorage( FACES_DATA, 0, CV_STORAGE_WRITE );

	// store all the data
	cvWriteInt( fileStorage, "nEigens", nEigens );
	cvWriteInt( fileStorage, "nTrainFaces", nTrainFaces );
	cvWrite(fileStorage, "trainPersonNumMat", personNumTruthMat, cvAttrList(0,0));
	cvWrite(fileStorage, "eigenValMat", eigenValMat, cvAttrList(0,0));
	cvWrite(fileStorage, "projectedTrainFaceMat", projectedTrainFaceMat, cvAttrList(0,0));
	cvWrite(fileStorage, "avgTrainImg", pAvgTrainImg, cvAttrList(0,0));
	for(i=0; i<nEigens; i++)
	{
		char varname[200];
		sprintf( varname, "eigenVect_%d", i );
		cvWrite(fileStorage, varname, eigenVectArr[i], cvAttrList(0,0));
	}

	// release the file-storage interface
	cvReleaseFileStorage( &fileStorage );
}
